import numpy as np
from collections import Counter

class KNN:
    def __init__(self, k=3):
        self.k = k  # 最近邻个数

    def fit(self, X_train, y_train):
        """训练模型，实际上只是存储训练数据"""
        self.X_train = X_train
        self.y_train = y_train

    def predict(self, X_test):
        """预测测试集标签"""
        return np.array([self._predict(x) for x in X_test])

    def _predict(self, x):
        """预测单个样本的标签"""
        # 计算与所有训练样本的距离
        distances = [self._euclidean_distance(x, x_train) for x_train in self.X_train]
        # 找到最近k个样本的索引
        k_indices = np.argsort(distances)[:self.k]
        # 获取最近k个样本的标签
        k_nearest_labels = [self.y_train[i] for i in k_indices]
        # 多数投票法确定标签
        most_common = Counter(k_nearest_labels).most_common(1)
        return most_common[0][0]

    def _euclidean_distance(self, x1, x2):
        """计算欧氏距离"""
        return np.sqrt(np.sum((x1 - x2) ** 2))

# ----------------- 测试 -----------------
if __name__ == "__main__":
    # 训练数据集
    X_train = np.array([[1, 2], [2, 3], [3, 3], [6, 5], [7, 8]])
    y_train = np.array([0, 0, 0, 1, 1])

    # 测试样本
    X_test = np.array([[1, 1], [7, 7]])

    # 创建KNN模型
    knn = KNN(k=3)
    knn.fit(X_train, y_train)

    # 预测
    predictions = knn.predict(X_test)
    print("预测结果:", predictions)
